#!/usr/bin/env python3
# -*- coding: UTF-8 -*-

import os
import sys
import time
import pandas
import threading
from netmiko import ConnectHandler

INFO_PATH = os.path.join(os.getcwd(), 'Sino-Bridge.xlsx')
LOCAL_TIME = time.strftime('%Y.%m.%d', time.localtime())
LOCK = threading.Lock()
POOL = threading.BoundedSemaphore(100)


def get_devices_info(info_file):
    """读取设备信息"""
    try:
        devices_dataframe = pandas.read_excel(info_file, sheet_name=0, dtype=str, keep_default_na=False)
    except FileNotFoundError:
        print(f'\n没有找到表格文件！\n')
        input('Press Enter to exit.')
        sys.exit(1)
    else:
        return devices_dataframe.to_dict('records')


def get_cmds_info(info_file):
    """读取每台设备的命令"""
    try:
        cmds_dataframe = pandas.read_excel(info_file, sheet_name=1, dtype=str, keep_default_na=False)
    except ValueError:
        print(f'\n表格文件缺失子表格信息！\n')
        input('Press Enter to exit.')
        sys.exit(1)
    else:
        # 变成 {host: [cmd1, cmd2, ...]} 结构
        cmds_dict = {}
        for _, row in cmds_dataframe.iterrows():
            host = row.get('host', '').strip()
            cmds = [str(c).strip() for c in row.values[1:] if str(c).strip()]
            cmds_dict[host] = cmds
        return cmds_dict


def inspection(login_info, cmds_list):
    """采集设备信息"""
    t_start = time.time()
    ssh = None

    try:
        ssh = ConnectHandler(**login_info)  # 不进入 enable 模式
    except Exception as ssh_error:
        with LOCK:
            err_msg = f'设备 {login_info.get("host")} 连接错误：{type(ssh_error).__name__}'
            print(err_msg)
            with open(os.path.join(os.getcwd(), LOCAL_TIME, 'err-log.log'), 'a', encoding='utf-8') as log:
                log.write(err_msg + '\n')
    else:
        log_path = os.path.join(os.getcwd(), LOCAL_TIME, f'{login_info["host"]}.log')
        with open(log_path, 'w', encoding='utf-8') as device_log_file:
            with LOCK:
                print(f'设备 {login_info["host"]} 正在采集...')

            for cmd in cmds_list:
                device_log_file.write('=' * 10 + ' ' + cmd + ' ' + '=' * 10 + '\n\n')
                try:
                    output = ssh.send_command(cmd, read_timeout=30)
                except Exception as e:
                    output = f'执行命令出错：{e}'
                device_log_file.write(output + '\n\n')

        with LOCK:
            print(f'设备 {login_info["host"]} 采集完成，用时 {round(time.time() - t_start, 1)} 秒。')

    finally:
        if ssh:
            ssh.disconnect()
        POOL.release()


if __name__ == '__main__':
    t1 = time.time()
    threading_list = []
    devices_info = get_devices_info(INFO_PATH)
    cmds_info = get_cmds_info(INFO_PATH)

    print(f'\n北京神州新桥科技有限公司西安分公司')
    print(f'\n采集开始...\n{"-" * 40}\n')

    if not os.path.exists(LOCAL_TIME):
        os.makedirs(LOCAL_TIME)
    else:
        try:
            os.remove(os.path.join(os.getcwd(), LOCAL_TIME, 'err-log.log'))
        except FileNotFoundError:
            pass

    for device_info in devices_info:
        host = device_info.get('host', '').strip()
        if not host:
            continue
        cmds_list = cmds_info.get(host, [])
        if not cmds_list:
            with LOCK:
                print(f'设备 {host} 未找到对应命令，跳过。')
            continue

        POOL.acquire()
        t = threading.Thread(target=inspection, args=(device_info, cmds_list))
        threading_list.append(t)
        t.start()

    for t in threading_list:
        t.join()

    try:
        with open(os.path.join(os.getcwd(), LOCAL_TIME, 'err-log.log'), 'r', encoding='utf-8') as log_file:
            file_lines = len(log_file.readlines())
    except FileNotFoundError:
        file_lines = 0

    t2 = time.time()
    print(f'\n{"-" * 40}\n')
    print(f'采集完成，共采集 {len(threading_list)} 台设备，{file_lines} 台异常，共用时 {round(t2 - t1, 1)} 秒。\n')
    input('Press Enter to exit.')

#不同的host执行不同的命令